{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GeerbrLA0uju"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\n",
        "Created by @[Adrish Dey](https://github.com/captain-pool) for [Google Summer of Code](https://summerofcode.withgoogle.com/) 2019"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yntM0JbY0uj5"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS, \n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UJeo2a5C0uj2"
      },
      "source": [
        "# Image Super Resolution using ESRGAN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ps4toA1d_tkc"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/image_enhancing\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/image_enhancing.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/hub/blob/master/examples/colab/image_enhancing.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/hub/examples/colab/image_enhancing.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://tfhub.dev/captain-pool/esrgan-tf2/1\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" /\u003eSee TF Hub model\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LkW9jAmt_zjB"
      },
      "source": [
        "This colab demonstrates use of TensorFlow Hub Module for Enhanced Super Resolution Generative Adversarial Network (*by Xintao Wang et.al.*) [[Paper](https://arxiv.org/pdf/1809.00219.pdf)] [[Code](https://github.com/captain-pool/GSOC/)]\n",
        "\n",
        "for image enhancing. *(Preferrably bicubically downsampled images).*\n",
        "\n",
        "Model trained on DIV2K Dataset (on bicubically downsampled images) on image patches of size 128 x 128."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LBGty4O_0ukJ"
      },
      "source": [
        "**Preparing Environment**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnyLTyUt0ukN"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import time\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import matplotlib.pyplot as plt\n",
        "os.environ[\"TFHUB_DOWNLOAD_PROGRESS\"] = \"True\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dremsFdh0ukX"
      },
      "outputs": [],
      "source": [
        "!wget \"https://user-images.githubusercontent.com/12981474/40157448-eff91f06-5953-11e8-9a37-f6b5693fa03f.png\" -O original.png"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DXot2kru0ukh"
      },
      "outputs": [],
      "source": [
        "# Declaring Constants\n",
        "IMAGE_PATH = \"original.png\"\n",
        "SAVED_MODEL_PATH = \"https://tfhub.dev/captain-pool/esrgan-tf2/1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KF_tHde-p3rn"
      },
      "source": [
        "**Defining Helper Functions**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IslbQmTj0ukz"
      },
      "outputs": [],
      "source": [
        "def preprocess_image(image_path):\n",
        "  \"\"\" Loads image from path and preprocesses to make it model ready\n",
        "      Args:\n",
        "        image_path: Path to the image file\n",
        "  \"\"\"\n",
        "  hr_image = tf.image.decode_image(tf.io.read_file(image_path))\n",
        "  # If PNG, remove the alpha channel. The model only supports\n",
        "  # images with 3 color channels.\n",
        "  if hr_image.shape[-1] == 4:\n",
        "    hr_image = hr_image[...,:-1]\n",
        "  hr_size = (tf.convert_to_tensor(hr_image.shape[:-1]) // 4) * 4\n",
        "  hr_image = tf.image.crop_to_bounding_box(hr_image, 0, 0, hr_size[0], hr_size[1])\n",
        "  hr_image = tf.cast(hr_image, tf.float32)\n",
        "  return tf.expand_dims(hr_image, 0)\n",
        "\n",
        "def save_image(image, filename):\n",
        "  \"\"\"\n",
        "    Saves unscaled Tensor Images.\n",
        "    Args:\n",
        "      image: 3D image tensor. [height, width, channels]\n",
        "      filename: Name of the file to save to.\n",
        "  \"\"\"\n",
        "  if not isinstance(image, Image.Image):\n",
        "    image = tf.clip_by_value(image, 0, 255)\n",
        "    image = Image.fromarray(tf.cast(image, tf.uint8).numpy())\n",
        "  image.save(\"%s.jpg\" % filename)\n",
        "  print(\"Saved as %s.jpg\" % filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uh1E2rBpnWxV"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "def plot_image(image, title=\"\"):\n",
        "  \"\"\"\n",
        "    Plots images from image tensors.\n",
        "    Args:\n",
        "      image: 3D image tensor. [height, width, channels].\n",
        "      title: Title to display in the plot.\n",
        "  \"\"\"\n",
        "  image = np.asarray(image)\n",
        "  image = tf.clip_by_value(image, 0, 255)\n",
        "  image = Image.fromarray(tf.cast(image, tf.uint8).numpy())\n",
        "  plt.imshow(image)\n",
        "  plt.axis(\"off\")\n",
        "  plt.title(title)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ycrCTvmlqBMD"
      },
      "source": [
        "#### Performing Super Resolution of images loaded from path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L7XpMk8Y0uk7"
      },
      "outputs": [],
      "source": [
        "hr_image = preprocess_image(IMAGE_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hWgCbUa_0ulG"
      },
      "outputs": [],
      "source": [
        "# Plotting Original Resolution image\n",
        "plot_image(tf.squeeze(hr_image), title=\"Original Image\")\n",
        "save_image(tf.squeeze(hr_image), filename=\"Original Image\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ouwEyKLA0ulO"
      },
      "outputs": [],
      "source": [
        "model = hub.load(SAVED_MODEL_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dz79ncnT0ulX"
      },
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "fake_image = model(hr_image)\n",
        "fake_image = tf.squeeze(fake_image)\n",
        "print(\"Time Taken: %f\" % (time.time() - start))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ABjkkJHC2jNL"
      },
      "outputs": [],
      "source": [
        "# Plotting Super Resolution Image\n",
        "plot_image(tf.squeeze(fake_image), title=\"Super Resolution\")\n",
        "save_image(tf.squeeze(fake_image), filename=\"Super Resolution\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tuKu18UYptkx"
      },
      "source": [
        "### Evaluating Performance of the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qdz55sxMgiwO"
      },
      "outputs": [],
      "source": [
        "!wget \"https://lh4.googleusercontent.com/-Anmw5df4gj0/AAAAAAAAAAI/AAAAAAAAAAc/6HxU8XFLnQE/photo.jpg64\" -O test.jpg\n",
        "IMAGE_PATH = \"test.jpg\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F6tMNtqy0ukq"
      },
      "outputs": [],
      "source": [
        "# Defining helper functions\n",
        "def downscale_image(image):\n",
        "  \"\"\"\n",
        "      Scales down images using bicubic downsampling.\n",
        "      Args:\n",
        "          image: 3D or 4D tensor of preprocessed image\n",
        "  \"\"\"\n",
        "  image_size = []\n",
        "  if len(image.shape) == 3:\n",
        "    image_size = [image.shape[1], image.shape[0]]\n",
        "  else:\n",
        "    raise ValueError(\"Dimension mismatch. Can work only on single image.\")\n",
        "\n",
        "  image = tf.squeeze(\n",
        "      tf.cast(\n",
        "          tf.clip_by_value(image, 0, 255), tf.uint8))\n",
        "\n",
        "  lr_image = np.asarray(\n",
        "    Image.fromarray(image.numpy())\n",
        "    .resize([image_size[0] // 4, image_size[1] // 4],\n",
        "              Image.BICUBIC))\n",
        "\n",
        "  lr_image = tf.expand_dims(lr_image, 0)\n",
        "  lr_image = tf.cast(lr_image, tf.float32)\n",
        "  return lr_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r2ANR1XDy77I"
      },
      "outputs": [],
      "source": [
        "hr_image = preprocess_image(IMAGE_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r_dautO6qbTV"
      },
      "outputs": [],
      "source": [
        "lr_image = downscale_image(tf.squeeze(hr_image))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jRw1x6xY0ulj",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# Plotting Low Resolution Image\n",
        "plot_image(tf.squeeze(lr_image), title=\"Low Resolution\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g--yyHg7qXCw"
      },
      "outputs": [],
      "source": [
        "model = hub.load(SAVED_MODEL_PATH)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZX-deZlhqaYz"
      },
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "fake_image = model(lr_image)\n",
        "fake_image = tf.squeeze(fake_image)\n",
        "print(\"Time Taken: %f\" % (time.time() - start))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AmSga6MSq1PB"
      },
      "outputs": [],
      "source": [
        "plot_image(tf.squeeze(fake_image), title=\"Super Resolution\")\n",
        "# Calculating PSNR wrt Original Image\n",
        "psnr = tf.image.psnr(\n",
        "    tf.clip_by_value(fake_image, 0, 255),\n",
        "    tf.clip_by_value(hr_image, 0, 255), max_val=255)\n",
        "print(\"PSNR Achieved: %f\" % psnr)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5YTBKCXPq9UZ"
      },
      "source": [
        "**Comparing Outputs size by side.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubdupldDypCy"
      },
      "outputs": [],
      "source": [
        "plt.rcParams['figure.figsize'] = [15, 10]\n",
        "fig, axes = plt.subplots(1, 3)\n",
        "fig.tight_layout()\n",
        "plt.subplot(131)\n",
        "plot_image(tf.squeeze(hr_image), title=\"Original\")\n",
        "plt.subplot(132)\n",
        "fig.tight_layout()\n",
        "plot_image(tf.squeeze(lr_image), \"x4 Bicubic\")\n",
        "plt.subplot(133)\n",
        "fig.tight_layout()\n",
        "plot_image(tf.squeeze(fake_image), \"Super Resolution\")\n",
        "plt.savefig(\"ESRGAN_DIV2K.jpg\", bbox_inches=\"tight\")\n",
        "print(\"PSNR: %f\" % psnr)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "image_enhancing.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.5.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
